# coding=utf-8
import traceback
import requests
from datetime import timedelta
from monthdelta import monthdelta
from pymongo import MongoClient
from simplemysql import SimpleMysql
from datetime import datetime
import Utils
from pandas import Series, DataFrame
import numpy as np
import sys

reload(sys)
sys.setdefaultencoding('utf-8')


import signal
import time




mgClient = MongoClient('121.40.54.235', 27017,
    username='root2',
    password='Dream2015',
    authSource="stockDataStore",
    authMechanism='SCRAM-SHA-1')
db = mgClient['stockDataStore']

mysqlDB = SimpleMysql(
    host="121.40.54.235",
    db="stock-grab",
    user="root",
    passwd="Dream2015",
    charset="utf8",
    autocommit=True,
    keep_alive=True  # try and reconnect timedout mysql connections?
)

def get_stock_period_indices_from_mongo(stock_code):
    stocks = db['stock_period_indices'].find({"code": stock_code},
                              projection={"period": 1,
                                          "net_profit": 1,
                                          "net_profit_opc": 1,
                                          "net_profit_koufei": 1,
                                          "npg_rate": 1,
                                          "total_assets": 1,
                                          "net_assets": 1,
                                          "net_assets_opc": 1,
                                          "total_shares": 1}, sort=[("period", 1)]);
    all_data = list(stocks)
    all_period = map(lambda x: x['period'], all_data);
    table_data = map(lambda x: [x['net_profit'], x['net_profit_opc'],x['npg_rate'],
                                x['total_assets'], x['net_assets'], x['net_assets_opc'], x['net_profit_koufei'],
                                x['total_shares']], all_data)
    df = DataFrame(table_data,
                   columns=['net_profit','net_profit_opc','npg_rate',
                            'total_assets', 'net_assets', 'net_assets_opc', 'net_profit_koufei',
                            'total_shares'], index=all_period)

    netvalue_in_period_df = Utils.calc_netvalue_in_period(df)
    df['net_profit_opc_in_period'] = netvalue_in_period_df['net_profit_opc']
    net_profit_opc_in_period_ttm = df['net_profit_opc_in_period'].rolling(4).sum()
    df['net_profit_opc_ttm'] = net_profit_opc_in_period_ttm

    df['net_profit_koufei_in_period'] = netvalue_in_period_df['net_profit_koufei']
    net_profit_koufei_in_period_ttm = df['net_profit_koufei_in_period'].rolling(4).sum()
    df['net_profit_koufei_ttm'] = net_profit_koufei_in_period_ttm

    return df


'''
获取day所在的季度财报数据，以此来计算PB、PE等
'''
def filter_period_indices(stock_period_df, day):
    all_period = stock_period_df.index
    for i in range(len(all_period)):
        period = all_period[i]
        if i == (len(all_period) - 1):
            '''
            比如现在是2018年7月1，目前只有1季度报表，半年报在2季度结束后2个月才有结果
            今天计算的财报数据还是以1季报为准，
            所以计算以哪个财报数据为准备？年报是3个月delay，半年报是2个月delay，季报以1个月delay
            '''
            delay = 0
            if period[4:6] == '03':
                delay = 5
            if period[4:6] == '09':
                delay = 6
            if period[4:6] == '06':
                delay = 4
            if period[4:6] == '12':
                delay = 4
            next_period = (datetime.strptime(period, '%Y%m%d') + monthdelta(delay)).strftime("%Y%m%d")
            if period <= day < next_period:
                return stock_period_df.iloc[i]
        else:
            next_period = all_period[i + 1]
            if period <= day < next_period:
                return stock_period_df.iloc[i]
    return None



def update_pe_koufei_of_stock(stock_code):
    period_indices_df = get_stock_period_indices_from_mongo(stock_code)
    all  = db['stock_daily_k'].find({'code': stock_code})
    for record in list(all):
        day = record[u'day'];
        period_indices_series_data_ = filter_period_indices(period_indices_df, day)
        if period_indices_series_data_ is not None:
            net_profit_ttm = period_indices_series_data_["net_profit_opc_ttm"]
            net_profit_koufei_ttm = period_indices_series_data_["net_profit_koufei_ttm"]
            mktcap = record[u'mktcap'];
            pe_ttm = float(0)
            pe_koufei_ttm = float(0)
            if not mktcap :
                continue
            if net_profit_ttm is not None and net_profit_ttm != 0:
                pe_ttm = float(mktcap) / net_profit_ttm
            if net_profit_koufei_ttm is not None and net_profit_koufei_ttm != 0:
                pe_koufei_ttm = float(mktcap) / net_profit_koufei_ttm

            if np.isnan(pe_ttm):
                pe_ttm = float(0)
            if np.isnan(pe_koufei_ttm):
                pe_koufei_ttm = float(0)

            print stock_code + "-" + record[u'day']
            print net_profit_ttm
            print net_profit_koufei_ttm
            print "origin:" + str(record[u'pe_ttm']) + " - " + str(record[u'pe_koufei_ttm'])
            print "new:" + str(pe_ttm) + " - " + str(pe_koufei_ttm) + "\n"
            db['stock_daily_k'].update_one({"_id": record[u'_id']},
                                     {'$set': {'pe': pe_ttm,
                                                'pe_koufei_ttm': pe_koufei_ttm,
                                               'code': record[u'code'],
                                               'modifyTime': datetime.now()}})


def update_stocks(stocks):
    for stock in list(stocks):
        if Utils.bypass_stock(stock['code']):
            Utils.logger.info(stock['code'] + " bypassed ")
            continue
        try:
            Utils.logger.info("--->start to repair pe/peg of " + stock["code"] + " "+ stock["name"] + "\n")
            update_pe_koufei_of_stock(stock["code"])
            Utils.logger.info("--->success to repair pe/peg of " + stock["code"] + " "+ stock["name"] + "\n")
        except:
            traceback.print_exc()
            Utils.logger.error("fail to collect daily k of " + stock['code'] + "\n")

def main():
    stocks = Utils.quick_access_cursor(db['stock'].find({"code": {"$regex": "^00"}}, projection={"code": 1, "name": 1}, sort=[("code", 1)]));
    update_stocks(stocks)
    stocks = Utils.quick_access_cursor(db['stock'].find({"code": {"$regex": "^30"}}, projection={"code": 1, "name": 1}, sort=[("code", 1)]));
    update_stocks(stocks)
    stocks = Utils.quick_access_cursor(db['stock'].find({"code": {"$regex": "^60"}}, projection={"code": 1, "name": 1}, sort=[("code", 1)]));
    update_stocks(stocks)

main()

#stocks=[{"code":'603898', 'name':'上海机场'}]
#update_stocks(stocks)

Utils.logger.info("CalcStockDailyK.py exit")
